import fasttext
from chat_bot.models.intent_fasttext import config
from chat_bot.models.intent_fasttext.utils import get_input

fasttext.FastText.eprint = lambda x: None


class FastTextModel(object):
    def __init__(self):
        self.model = fasttext.load_model(config.model_path)
        pass

    def predict(self, query):
        input = get_input(query)
        label, percent = self.model.predict(input)
        return {"label": label[0].replace('__label__', ''), "confidence": percent[0]}


if __name__ == '__main__':
    _fasttext = FastTextModel()

    # 预测
    print(_fasttext.predict('丹参是啥？'))
    print(_fasttext.predict("高粱紫轮病有什么物理防治手段？"))
    print(_fasttext.predict("梅树白囊蓑蛾的生活习性是啥？"))
    print(_fasttext.predict("蚜虫有触角吗"))
    print(_fasttext.predict("蚜虫的颜色?"))
    print(_fasttext.predict("荞麦茶多少钱一斤? "))
    print(_fasttext.predict("玉米的花有多少朵"))
    print(_fasttext.predict("培育玉米的所需光照强度"))
    print(_fasttext.predict("培育玉米的所需的温度"))
    print(_fasttext.predict("种植玉米土壤的PH值"))
    print(_fasttext.predict("培育水稻是不是很难啊"))
    print(_fasttext.predict("水稻的繁殖方式"))
    print(_fasttext.predict("异狄氏剂使用要兑多少水"))
    print(_fasttext.predict("异狄氏剂主要对付什么害虫"))
    print(_fasttext.predict("异狄氏剂对感染了霜霉病的葡萄有作用吗"))
